<!DOCTYPE html>
<html>
<head>
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@0.12.0"> </script>
  <script src="../node_modules/file-saver/FileSaver.js"></script>
</head>
<body>

  <div class="row side-by-side">
    <button
      class="waves-effect waves-light btn"
      onclick="save()"
    >
      Save
    </button>
    <button
      class="waves-effect waves-light btn"
      onclick="save(true)"
    >
      Save Tiny
    </button>
  </div>

  <script>
    function toDataArray(tensor) {
      return Array.from(tensor.dataSync())
    }

    function flatten(arrs) {
      return arrs.reduce((flat, arr) => flat.concat(arr))
    }

    function initWeights(initializer, isTiny) {
      function initConvWeights(inChannels, outChannels) {
        return flatten(
          [
            // filter
            initializer.apply([3, 3, inChannels, outChannels]),
            // bias
            tf.zeros([outChannels])
          ]
            .map(toDataArray)
        )
      }

      function initSeparableConvWeights(inChannels, outChannels) {
        return flatten(
          [
            // depthwise filter
            initializer.apply([3, 3, inChannels, 1]),
            // pointwise filter
            initializer.apply([1, 1, inChannels, outChannels]),
            // bias
            tf.zeros([outChannels])
          ]
            .map(toDataArray)
        )
      }

      const separableConvWeights = isTiny
        ? flatten([
          initConvWeights(3, 32),
          initSeparableConvWeights(32, 32),
          initSeparableConvWeights(32, 32),
          initSeparableConvWeights(32, 64),
          initSeparableConvWeights(64, 64),
          initSeparableConvWeights(64, 64),
          initSeparableConvWeights(64, 128),
          initSeparableConvWeights(128, 128),
          initSeparableConvWeights(128, 128)
        ])
        : flatten([
          initConvWeights(3, 32),
          initSeparableConvWeights(32, 32),
          initSeparableConvWeights(32, 32),
          initSeparableConvWeights(32, 32),
          initSeparableConvWeights(32, 64),
          initSeparableConvWeights(64, 64),
          initSeparableConvWeights(64, 64),
          initSeparableConvWeights(64, 64),
          initSeparableConvWeights(64, 128),
          initSeparableConvWeights(128, 128),
          initSeparableConvWeights(128, 128),
          initSeparableConvWeights(128, 128),
          initSeparableConvWeights(128, 256),
          initSeparableConvWeights(256, 256),
          initSeparableConvWeights(256, 256),
          initSeparableConvWeights(256, 256)
        ])

      const fc = flatten(
        [
          initializer.apply([1, 1, isTiny ? 128 : 256, 136]),
          // bias
          tf.zeros([136])
        ]
          .map(toDataArray)
      )
      return new Float32Array(separableConvWeights.concat(fc))
    }

    function save(isTiny = false) {
      const initialWeights = initWeights(
        tf.initializers.glorotNormal(),
        isTiny
      )
      saveAs(new Blob([initialWeights]), `initial_glorot.weights`)
    }

  </script>
</body>
</html>